#pragma once

// Extended MultiSelectTree styles
static const DWORD MTVS_EX_NOMARQUEE = 0x00000001;

// New control notifications
static const UINT TVN_ITEMSELECTING = 0x0001;
static const UINT TVN_ITEMSELECTED = 0x0002;

static bool operator==(const CTreeItem& ti1, const CTreeItem& ti2)
{
	return ti1.m_hTreeItem == ti2.m_hTreeItem;
}

class CMultiSelectTreeViewCtrl :
	public CWindowImpl<CMultiSelectTreeViewCtrl, CTreeViewCtrlEx, CWinTraitsOR<TVS_SHOWSELALWAYS>>,
	public CCustomDraw<CMultiSelectTreeViewCtrl>
{
public:
	DECLARE_WND_SUPERCLASS(_T("WTL_MultiSelectTree"), GetWndClassName())

	DWORD m_dwExStyle;              // Additional styles
	CTreeItem m_hExtSelStart;       // Item where SHIFT was last pressed
	bool m_bMarquee;                // Are we drawing rubberband?
	CPoint m_ptDragStart;            // Point where rubberband started
	CPoint m_ptDragOld;              // Last mousepos of rubberband

	CSimpleMap<HTREEITEM, bool> m_aData;

	CMultiSelectTreeViewCtrl() :m_dwExStyle(0), m_bMarquee(false) {

	}

	CTreeItem _MakeItem(HTREEITEM hItem) const
	{
		return CTreeItem(hItem, (CTreeViewCtrlEx*)this);
	}

	BOOL SubclassWindow(HWND hWnd)
	{
		ATLASSERT(m_hWnd == NULL);
		ATLASSERT(::IsWindow(hWnd));

		BOOL bRet = CWindowImpl<CMultiSelectTreeViewCtrl, CTreeViewCtrlEx, CWinTraitsOR<TVS_SHOWSELALWAYS>>::SubclassWindow(hWnd);
		if (bRet)
			_Init();
		return bRet;
	}

	BOOL SelectItem(HTREEITEM hItem, BOOL bSelect)
	{
		ATLASSERT(::IsWindow(m_hWnd));

		_SelectItem(hItem, bSelect == TRUE);
		if (bSelect)
			CWindowImpl<CMultiSelectTreeViewCtrl, CTreeViewCtrlEx, CWinTraitsOR<TVS_SHOWSELALWAYS>>::SelectItem(hItem);
		return TRUE;
	}

	BOOL SelectAllItems(BOOL bSelect)
	{
		ATLASSERT(::IsWindow(m_hWnd));

		for (int i = 0; i < m_aData.GetSize(); i++)
		{
			_SelectItem(i, bSelect == TRUE);
		}
		return TRUE;
	}

	BOOL IsItemSelected(HTREEITEM hItem)
	{
		ATLASSERT(::IsWindow(m_hWnd));

		int iIndex = m_aData.FindKey(hItem);
		if (iIndex >= 0)
		{
			return m_aData.GetValueAt(iIndex) ? TRUE : FALSE;
		}
		return FALSE;
	}

	CTreeItem GetSelectedItem() const
	{
		ATLASSERT(false);  // Not usable!
		return GetFirstSelectedItem();
	}

	CTreeItem GetFocusItem() const
	{
		return CWindowImpl<CMultiSelectTreeViewCtrl, CTreeViewCtrlEx, CWinTraitsOR<TVS_SHOWSELALWAYS>>::GetSelectedItem();
	}

	UINT GetItemState(HTREEITEM hItem, UINT nStateMask) const
	{
		UINT nRes = CWindowImpl<CMultiSelectTreeViewCtrl, CTreeViewCtrlEx, CWinTraitsOR<TVS_SHOWSELALWAYS>>::GetItemState(hItem, nStateMask);
		if ((nStateMask & TVIS_SELECTED) != 0)
		{
			int iIndex = m_aData.FindKey(hItem);
			if (iIndex >= 0)
			{
				nRes &= ~TVIS_SELECTED;
				if (m_aData.GetValueAt(iIndex))
					nRes |= TVIS_SELECTED;
			}
		}
		return nRes;
	}

	CTreeItem GetFirstSelectedItem() const
	{
		HTREEITEM item = NULL;
		if (m_aData.GetSize() > 0)
		{
			for (int i = 0; i < m_aData.GetSize(); i++)
			{
				if (m_aData.GetValueAt(i))
				{
					item = m_aData.GetKeyAt(i);
					break;
				}
			}
		}
		return _MakeItem(item);
	}

	CTreeItem GetNextSelectedItem(HTREEITEM hItem) const
	{
		HTREEITEM item = NULL;
		int iIndex = m_aData.FindKey(hItem);
		if (iIndex >= 0)
		{
			for (int i = iIndex + 1; i < m_aData.GetSize(); i++)
			{
				if (m_aData.GetValueAt(i))
				{
					item = m_aData.GetKeyAt(i);
					break;
				}
			}
		}
		return _MakeItem(item);
	}

	int GetSelectedCount() const
	{
		int nCount = 0;
		for (int i = 0; i < m_aData.GetSize(); i++)
		{
			if (m_aData.GetValueAt(i))
				nCount++;
		}
		return nCount;
	}

	// Implementation

	void _Init()
	{
		ATLASSERT(::IsWindow(m_hWnd));

		ModifyStyle(TVS_SHOWSELALWAYS, 0);
	}

	void _SelectItem(int iIndex, bool bSelect, int action = TVC_UNKNOWN)
	{
		if (iIndex < 0)
			return;
		bool bSelected = m_aData.GetValueAt(iIndex);
		// Don't change if state is already updated (avoids flicker)
		if (bSelected == bSelect)
			return;
		HTREEITEM hItem = m_aData.GetKeyAt(iIndex);
		CWindow parent = GetParent();
		// Send notifications
		NMTREEVIEW nmtv = { 0 };
		nmtv.hdr.code = TVN_ITEMSELECTING;
		nmtv.hdr.hwndFrom = m_hWnd;
		nmtv.hdr.idFrom = GetDlgCtrlID();
		nmtv.action = action;
		nmtv.itemNew.hItem = hItem;
		nmtv.itemNew.lParam = GetItemData(hItem);
		nmtv.itemNew.state = bSelect ? TVIS_SELECTED : 0;
		nmtv.itemNew.stateMask = TVIS_SELECTED;
		if (parent.SendMessage(WM_NOTIFY, nmtv.hdr.idFrom, (LPARAM)&nmtv) != 0)
			return;
		// Change state
		m_aData.SetAtIndex(iIndex, hItem, bSelect);
		// Repaint item
		CRect rcItem;
		if (GetItemRect(hItem, &rcItem, FALSE))
			InvalidateRect(&rcItem, TRUE);
		// More notifications
		nmtv.hdr.code = TVN_ITEMSELECTED;
		parent.SendMessage(WM_NOTIFY, nmtv.hdr.idFrom, (LPARAM)&nmtv);
	}

	void _SelectItem(HTREEITEM hItem, bool bSelect, int action = TVC_UNKNOWN)
	{
		_SelectItem(m_aData.FindKey(hItem), bSelect, action);
	}

	void _SelectTree(HTREEITEM hItem, HTREEITEM hGoal, int action)
	{
		if (!_SelectTreeSub(hItem, hGoal, action))
			return;
		hItem = GetParentItem(hItem);
		while ((hItem = GetNextSiblingItem(hItem)) != NULL)
		{
			if (!_SelectTreeSub(hItem, hGoal, action))
				return;
		}
	}

	bool _SelectTreeSub(HTREEITEM hItem, HTREEITEM hGoal, int action)
	{
		while (hItem != NULL)
		{
			_SelectItem(hItem, true, action);
			if (hItem == hGoal)
				return false;
			if ((CWindowImpl<CMultiSelectTreeViewCtrl, CTreeViewCtrlEx, CWinTraitsOR<TVS_SHOWSELALWAYS>>::GetItemState(hItem, TVIS_EXPANDED) & TVIS_EXPANDED) != 0)
			{
				if (!_SelectTreeSub(GetChildItem(hItem), hGoal, action))
					return false;
			}
			hItem = GetNextSiblingItem(hItem);
		}
		return true;
	}

	void _SelectBox(CRect rc)
	{
		HTREEITEM hItem = GetFirstVisibleItem();
		while (hItem != NULL)
		{
			int i = m_aData.FindKey(hItem);
			if (i >= 0 && !m_aData.GetValueAt(i)) // ignore already selected
			{
				CRect rcItem, rcTemp;
				GetItemRect(hItem, &rcItem, TRUE);
				_SelectItem(hItem, rcTemp.IntersectRect(&rcItem, &rc) == TRUE, TVC_BYMOUSE);
			}
			hItem = GetNextVisibleItem(hItem);
		}
	}

	void _DrawDragRect(CPoint pt)
	{
		CClientDC dc(m_hWnd);
		CSize szFrame(1, 1);
		CRect rect(m_ptDragStart, pt);
		rect.NormalizeRect();
		//CRect rectOld(m_ptDragStart, m_ptDragOld);
		//rectOld.NormalizeRect();
		dc.DrawDragRect(&rect, szFrame, NULL/*&rectOld*/, szFrame);
	}

	BEGIN_MSG_MAP_EX(CMultiSelectTreeViewImpl)
		MESSAGE_HANDLER_EX(WM_CREATE, OnCreate)
		MSG_WM_DESTROY(OnDestroy)
		MSG_WM_KEYDOWN(OnKeyDown)
		MSG_WM_KEYUP(OnKeyUp)
		MSG_WM_CHAR(OnChar)
		MSG_WM_SETFOCUS(OnSetFocus)
		MSG_WM_LBUTTONDOWN(OnLButtonDown)
		MSG_WM_LBUTTONUP(OnLButtonUp)
		MSG_WM_MOUSEMOVE(OnMouseMove)
		MSG_WM_CAPTURECHANGED(OnCaptureChanged)
		MESSAGE_HANDLER_EX(TVM_INSERTITEM, OnInsertItem)
		REFLECTED_NOTIFY_CODE_HANDLER_EX(TVN_DELETEITEM, OnDeleteItem)
		CHAIN_MSG_MAP_ALT(CCustomDraw<CMultiSelectTreeViewCtrl>, 1)
	END_MSG_MAP()

	LRESULT OnCreate(UINT uMsg, WPARAM wParam, LPARAM lParam)
	{
		LRESULT lRes = DefWindowProc();
		_Init();
		return lRes;
	}

	void OnDestroy()
	{
		m_aData.RemoveAll();
		SetMsgHandled(FALSE);
	}

	void OnKeyDown(UINT nChar, UINT nRepCnt, UINT nFlags)
	{
		if (nChar == VK_SHIFT)
			m_hExtSelStart = GetFocusItem();

		if (::GetKeyState(VK_SHIFT) < 0 && m_hExtSelStart == GetFocusItem())
		{
			switch (nChar)
			{
				case VK_UP:
				case VK_DOWN:
				case VK_HOME:
				case VK_END:
				case VK_NEXT:
				case VK_PRIOR:
					for (int i = 0; i < m_aData.GetSize(); i++)
					{
						_SelectItem(i, false, TVC_BYKEYBOARD);
					}
			}
		}
		SetMsgHandled(FALSE);
	}

	void OnKeyUp(UINT nChar, UINT nRepCnt, UINT nFlags)
	{
		if (::GetKeyState(VK_SHIFT) < 0)
		{
			switch (nChar)
			{
				case VK_UP:
				case VK_DOWN:
				case VK_HOME:
				case VK_END:
				case VK_NEXT:
				case VK_PRIOR:
					HTREEITEM hItem = GetFocusItem();
					// Is current or first-shift-item the upper item?
					CRect rcItem1, rcItem2;
					GetItemRect(m_hExtSelStart, &rcItem1, TRUE);
					GetItemRect(hItem, &rcItem2, TRUE);
					// Select from current item to item where SHIFT was pressed
					if (rcItem1.top > rcItem2.top)
						_SelectTree(hItem, m_hExtSelStart, TVC_BYKEYBOARD);
					else
						_SelectTree(m_hExtSelStart, hItem, TVC_BYKEYBOARD);
					_SelectItem(hItem, true, TVC_BYKEYBOARD);
			}
		}
		SetMsgHandled(FALSE);
	}

	void OnChar(UINT nChar, UINT nRepCnt, UINT nFlags)
	{
		if (nChar == VK_SPACE)
		{
			HTREEITEM hItem = GetFocusItem();
			_SelectItem(hItem, IsItemSelected(hItem) == TRUE/*GetItemState(hItem, TVIS_SELECTED) & TVIS_SELECTED) == 0*/, TVC_BYKEYBOARD);
			return;
		}
		SetMsgHandled(FALSE);
	}

	void OnSetFocus(CWindow wndOld)
	{
		DefWindowProc();
		// FIX: We really need the focus-rectangle in this control since it
		//      improves the navigation a lot. So let's ask Windows to display it.
		SendMessage(WM_UPDATEUISTATE, MAKEWPARAM(UIS_CLEAR, UISF_HIDEFOCUS));
	}

	void OnLButtonDown(UINT nFlags, CPoint point)
	{
		SetMsgHandled(FALSE);

		// Hit-test and figure out where we're clicking...
		TVHITTESTINFO hti = { 0 };
		hti.pt = point;
		HTREEITEM hItem = HitTest(&hti);
		if ((hItem == NULL || (hti.flags & TVHT_ONITEMRIGHT) != 0))
		{
			if ((m_dwExStyle & MTVS_EX_NOMARQUEE) == 0 && ::DragDetect(m_hWnd, point))
			{
				// Great we're dragging a rubber-band
				// Clear selection if CTRL is not down
				if (::GetKeyState(VK_CONTROL) >= 0)
				{
					for (int i = 0; i < m_aData.GetSize(); i++)
					{
						_SelectItem(i, false, TVC_BYMOUSE);
					}
					UpdateWindow();
				}
				// Now start drawing the rubber-band...
				SetCapture();
				m_ptDragStart = m_ptDragOld = point;
				_DrawDragRect(point);
				m_bMarquee = true;
				SetMsgHandled(FALSE); //SetMsgHandled(TRUE);
				return;
			}
		}

		if (hItem == NULL)
			return;

		if ((hti.flags & TVHT_ONITEMBUTTON) != 0)
			return;

		// Great, let's do an advanced selection
		if ((hti.flags & TVHT_ONITEMRIGHT) != 0)
		{
			for (int i = 0; i < m_aData.GetSize(); i++)
			{
				_SelectItem(i, false, TVC_BYMOUSE);
			}
			return;
		}
		int iIndex = m_aData.FindKey(hItem);
		if (iIndex < 0)
			return;
		// Simulate drag'n'drop?
		if (m_aData.GetValueAt(iIndex) && (GetStyle() & TVS_DISABLEDRAGDROP) == 0 && ::DragDetect(m_hWnd, point))
		{
			NMTREEVIEW nmtv = { 0 };
			nmtv.hdr.code = TVN_BEGINDRAG;
			nmtv.hdr.hwndFrom = m_hWnd;
			nmtv.hdr.idFrom = GetDlgCtrlID();
			nmtv.itemNew.hItem = hItem;
			nmtv.itemNew.lParam = GetItemData(hItem);
			CWindow parent = GetParent();
			parent.SendMessage(WM_NOTIFY, nmtv.hdr.idFrom, (LPARAM)&nmtv);
		}
		bool bSelected = m_aData.GetValueAt(iIndex);
		if (::GetKeyState(VK_SHIFT) < 0)
		{
			// Is current or first-shift-item the upper item?
			CRect rcItem1, rcItem2;
			GetItemRect(m_hExtSelStart, &rcItem1, TRUE);
			GetItemRect(hItem, &rcItem2, TRUE);
			// Select from current item to item where SHIFT was pressed
			if (rcItem1.top > rcItem2.top)
				_SelectTree(hItem, m_hExtSelStart, TVC_BYMOUSE);
			else
				_SelectTree(m_hExtSelStart, hItem, TVC_BYMOUSE);
		}
		else if (::GetKeyState(VK_CONTROL) < 0)
		{
			// Just toggle item
			_SelectItem(iIndex, !bSelected, TVC_BYMOUSE);
		}
		else
		{
			// Remove current selection and replace it with clicked item
			for (int i = 0; i < m_aData.GetSize(); i++)
			{
				_SelectItem(i, i == iIndex, TVC_BYMOUSE);
			}
		}
	}

	void OnLButtonUp(UINT nFlags, CPoint point)
	{
		if (m_bMarquee)
			ReleaseCapture();
		SetMsgHandled(FALSE);
	}

	void OnMouseMove(UINT nFlags, CPoint point)
	{
		if (m_bMarquee)
		{
			CRect rc(m_ptDragStart, point);
			_DrawDragRect(m_ptDragOld);
			rc.NormalizeRect();
			_SelectBox(rc);
			UpdateWindow();
			_DrawDragRect(point);
			m_ptDragOld = point;
		}
		SetMsgHandled(FALSE);
	}

	void OnCaptureChanged(CWindow wnd)
	{
		if (m_bMarquee)
		{
			_DrawDragRect(m_ptDragOld);
			m_bMarquee = false;
		}
		SetMsgHandled(FALSE);
	}

	LRESULT OnInsertItem(UINT uMsg, WPARAM wParam, LPARAM lParam)
	{
		HTREEITEM hItem = (HTREEITEM)DefWindowProc(uMsg, wParam, lParam);
		if (hItem == NULL)
			return (LRESULT)hItem;
		// We manage a bit of extra information for each item. We'll store
		// this in an ATL::CSimpleMap. Not a particular speedy structure for lookups.
		// Don't keep too many items around in the tree!
		m_aData.Add(hItem, false);
		return (LRESULT)hItem;
	}

	LRESULT OnDeleteItem(NMHDR* pnmh)
	{
		const NMTREEVIEW* lpNMTV = (NMTREEVIEW*)pnmh;
		m_aData.Remove(lpNMTV->itemNew.hItem);
		return 0;
	}

	// Custom Draw

	DWORD OnPrePaint(int /*idCtrl*/, NMCUSTOMDRAW* /*lpNMCustomDraw*/)
	{
		return CDRF_NOTIFYITEMDRAW;   // We need per-item notifications
	}

	DWORD OnItemPrePaint(int /*idCtrl*/, NMCUSTOMDRAW* lpNMCustomDraw)
	{
		NMTVCUSTOMDRAW* lpTVCD = (NMTVCUSTOMDRAW*)lpNMCustomDraw;
		HTREEITEM hItem = (HTREEITEM)lpTVCD->nmcd.dwItemSpec;
		int iIndex = m_aData.FindKey(hItem);
		if (iIndex >= 0)
		{
			bool bSelected = m_aData.GetValueAt(iIndex);
			// Trick TreeView into displaying correct selection colors
			if (bSelected)
			{
				lpTVCD->clrText = ::GetSysColor(COLOR_HIGHLIGHTTEXT);
				lpTVCD->clrTextBk = ::GetSysColor(COLOR_HIGHLIGHT);
			}
			else
			{
				// Special case of tree-item actually have selection, but our
				// state says it is currently not selected (CTRL+click on same item twice).
				if ((lpTVCD->nmcd.uItemState & CDIS_SELECTED) != 0)
				{
					COLORREF clrText = GetTextColor();
					if (clrText == CLR_NONE)
						clrText = ::GetSysColor(COLOR_WINDOWTEXT);
					COLORREF clrBack = GetBkColor();
					if (clrBack == CLR_NONE)
						clrBack = ::GetSysColor(COLOR_WINDOW);
					//CDCHandle dc = lpTVCD->nmcd.hdc;
					//dc.SetTextColor(clrText);
					//dc.SetBkColor(clrBack);
					lpTVCD->clrText = clrText;
					lpTVCD->clrTextBk = clrBack;
				}
			}
			return CDRF_NEWFONT;
		}
		return CDRF_DODEFAULT;
	}

};